#!/usr/bin/env python3
"""
数据库初始化和迁移脚本
确保表结构正确并添加必要的索引
"""
import os
import sys
import json
import logging
from datetime import datetime

# 添加项目根目录到Python路径
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), '../..')))

from backend.common.config import config_manager
from backend.database.db import db_manager

logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger('database_migration')

def migrate_database():
    """执行数据库迁移"""
    try:
        logger.info(f"开始数据库迁移，数据库类型: {db_manager.db_type}")
        
        # 检查并创建必要的表和索引
        with db_manager.get_cursor() as cursor:
            # 创建users表（如果不存在）
            if db_manager.db_type == 'sqlite':
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id VARCHAR(255) PRIMARY KEY,
                    username TEXT UNIQUE NOT NULL,
                    email TEXT UNIQUE,
                    password_hash TEXT NOT NULL,
                    nickname TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            else:
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS users (
                    id VARCHAR(255) PRIMARY KEY,
                    username VARCHAR(255) UNIQUE NOT NULL,
                    email VARCHAR(255) UNIQUE,
                    password_hash TEXT NOT NULL,
                    nickname VARCHAR(255),
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            db_manager._ensure_user_identifiers(cursor)

            # 创建sessions表（如果不存在）
            if db_manager.db_type == 'sqlite':
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    id VARCHAR(255) PRIMARY KEY,
                    user_id VARCHAR(255) DEFAULT 'default',
                    title TEXT NOT NULL DEFAULT '新会话',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            else:
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS sessions (
                    id VARCHAR(255) PRIMARY KEY,
                    user_id VARCHAR(255) DEFAULT 'default',
                    title TEXT NOT NULL DEFAULT '新会话',
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            
            # 创建sessions表的索引
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_sessions_user_id ON sessions(user_id)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_sessions_created_at ON sessions(created_at)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_sessions_updated_at ON sessions(updated_at)")
            
            # 创建messages表（如果不存在）
            if db_manager.db_type == 'sqlite':
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    session_id VARCHAR(255) NOT NULL,
                    role VARCHAR(50) NOT NULL,
                    content TEXT NOT NULL,
                    meta TEXT,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            else:
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS messages (
                    id SERIAL PRIMARY KEY,
                    session_id VARCHAR(255) NOT NULL,
                    role VARCHAR(50) NOT NULL,
                    content TEXT NOT NULL,
                    meta JSONB,
                    created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            
            # 创建messages表的索引
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_session_id ON messages(session_id)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_created_at ON messages(created_at)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_role ON messages(role)")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_session_role ON messages(session_id, role)")

            # 确保messages表存在user_id列
            if db_manager.db_type == 'sqlite':
                cursor.execute("PRAGMA table_info(messages)")
                columns = {row['name'] if isinstance(row, dict) else row[1] for row in cursor.fetchall()}
                if 'user_id' not in columns:
                    cursor.execute("ALTER TABLE messages ADD COLUMN user_id VARCHAR(255) DEFAULT 'default'")
            else:
                cursor.execute("ALTER TABLE messages ADD COLUMN IF NOT EXISTS user_id VARCHAR(255) DEFAULT 'default'")
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_messages_user_id ON messages(user_id)")
            
            # 创建会话模型映射表
            if db_manager.db_type == 'sqlite':
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS session_model_map (
                    session_id VARCHAR(255) PRIMARY KEY,
                    model VARCHAR(255) NOT NULL,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            else:
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS session_model_map (
                    session_id VARCHAR(255) PRIMARY KEY,
                    model VARCHAR(255) NOT NULL,
                    updated_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP
                )
                """)
            cursor.execute("CREATE INDEX IF NOT EXISTS idx_session_model_updated_at ON session_model_map(updated_at)")
            
            
            # 添加迁移记录
            migration_version = "1.0.0"
            migration_time = datetime.utcnow().isoformat()
            
            # 创建迁移历史表（如果不存在）
            if db_manager.db_type == 'sqlite':
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS migrations (
                    id INTEGER PRIMARY KEY AUTOINCREMENT,
                    version TEXT NOT NULL,
                    applied_at TEXT NOT NULL,
                    details TEXT
                )
                """)
            else:
                cursor.execute("""
                CREATE TABLE IF NOT EXISTS migrations (
                    id SERIAL PRIMARY KEY,
                    version TEXT NOT NULL,
                    applied_at TEXT NOT NULL,
                    details JSONB
                )
                """)
            
            # 记录本次迁移
            migration_details = {
                "tables_created": ["users", "sessions", "messages", "session_model_map"],
                "indexes_created": [
                    "idx_users_username",
                    "idx_users_email",
                    "idx_sessions_user_id", "idx_sessions_created_at", "idx_sessions_updated_at",
                    "idx_messages_session_id", "idx_messages_created_at", "idx_messages_role",
                    "idx_messages_session_role", "idx_messages_user_id",
                    "idx_session_model_updated_at"
                ],
                "db_type": db_manager.db_type
            }
            
            if db_manager.db_type == 'sqlite':
                cursor.execute(
                    "INSERT INTO migrations (version, applied_at, details) VALUES (?, ?, ?)",
                    (migration_version, migration_time, json.dumps(migration_details))
                )
            else:
                cursor.execute(
                    "INSERT INTO migrations (version, applied_at, details) VALUES (%s, %s, %s)",
                    (migration_version, migration_time, migration_details)
                )
            
            db_manager.connection.commit()
            logger.info(f"数据库迁移成功完成，版本: {migration_version}")
            
            # 验证表结构
            if db_manager.db_type == 'sqlite':
                cursor.execute("SELECT name FROM sqlite_master WHERE type='table'")
            else:
                cursor.execute("SELECT table_name FROM information_schema.tables WHERE table_schema='public'")
            tables = cursor.fetchall()
            logger.info(f"数据库中的表: {[table['name'] if isinstance(table, dict) else table[0] for table in tables]}")
            
            return True
    except Exception as e:
        logger.error(f"数据库迁移失败: {str(e)}")
        if hasattr(db_manager.connection, 'rollback'):
            db_manager.connection.rollback()
        return False

def verify_schema():
    """验证数据库模式"""
    try:
        logger.info("验证数据库模式")
        
        with db_manager.get_cursor() as cursor:
            # 检查sessions表
            if db_manager.db_type == 'sqlite':
                cursor.execute("PRAGMA table_info(sessions)")
            else:
                cursor.execute(
                    "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = 'sessions'"
                )
            sessions_columns = cursor.fetchall()
            logger.info(f"sessions表列: {sessions_columns}")
            
            # 检查messages表
            if db_manager.db_type == 'sqlite':
                cursor.execute("PRAGMA table_info(messages)")
            else:
                cursor.execute(
                    "SELECT column_name, data_type FROM information_schema.columns WHERE table_name = 'messages'"
                )
            messages_columns = cursor.fetchall()
            logger.info(f"messages表列: {messages_columns}")
            
            return True
    except Exception as e:
        logger.error(f"数据库模式验证失败: {str(e)}")
        return False

if __name__ == "__main__":
    logger.info("启动数据库迁移工具")
    
    # 初始化数据库连接
    if not db_manager.connection:
        db_manager._connect()
    
    # 执行迁移
    if migrate_database():
        verify_schema()
        logger.info("数据库迁移工具执行完成")
        sys.exit(0)
    else:
        logger.error("数据库迁移工具执行失败")
        sys.exit(1)
